import gradio as gr
import numpy as np
import joblib
import cv2

# 加载模型
with open('best_model.pkl', 'rb') as f:
    model = joblib.load(f)  # 使用joblib.load代替pickle.load

def predict_image(image):
    if len(image.shape) == 3:
        img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        img = image
        
    img = cv2.resize(img, (32, 16))
    img = img / 255.0  # 归一化
    # 添加这些打印语句来检查数据
    print("输入图像形状:", img.shape)
    
    
    
    image_flattened = img.flatten().reshape(1, -1)
    prediction = model.predict(image_flattened)[0]
    
    # 打印原始预测结果
    print("模型预测结果:", prediction)
    
    return {
        "猫": float(prediction == 0),
        "狗": float(prediction == 1)
    }

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(),  # 移除 type="file"
    outputs=gr.Label(num_top_classes=2)
)

# 启动应用
iface.launch()